#include "PowerSpectrum.h"

#include <iostream>
#include <fstream>
using std::endl;
using std::cout;
#include <string>
#include <sstream>
using std::string;
#include <vector>
using std::vector;

#include "dbg.h"
#include "TMV.h"
using tmv::Vector;
using tmv::Matrix;

#define DATECODE "0913"

std::ostream* dbgout=0;
bool XDEBUG=false;

const double PI = 3.141592653589793;
const double ARCMIN = PI/180./60.;

const double map_minr = 1.0*ARCMIN;
const double map_maxr = 40.0*ARCMIN;
const double var_minr = 40.0*ARCMIN;
const double var_maxr = 100.0*ARCMIN;

const double pred_map_minr = 1.0*ARCMIN;
const double pred_map_maxr = 40.0*ARCMIN;
const int npred_map = 100;
const double pred_var_minr = 1.0*ARCMIN;
const double pred_var_maxr = 100.0*ARCMIN;
const int npred_var = 100;

const double singthresh = 0.01;
const int datafrac = 1;

void ReadData(Vector<double>*& r, Vector<double>*& mean, Matrix<double>*& cov, 
    size_t& nmapout)
{
  dbg<<"Start ReadData\n";

  const size_t nfields = 12;
  const string datafile = "/home/mjarvis/plots/m2.out.";
  const char* fieldcodes[nfields] = { 
    "A", "G", "H", "J", "K", "L", "M", "N", "Q", "R", "T", "X" };

  vector<vector<double> > rmap(nfields);
  vector<vector<double> > map(nfields);
  //vector<vector<double> > sigmap(nfields);
  vector<vector<double> > rvar(nfields);
  vector<vector<double> > var(nfields);
  //vector<vector<double> > sigvar(nfields);

  for(size_t k=0;k<nfields;k++) {
    string name = datafile + DATECODE + '.' + fieldcodes[k];
    std::ifstream fin(name.c_str());
    if (!fin) myerror("opening data file:  ",name.c_str());

    dbg<<"Reading file: "<<name<<endl;
    string line;
    int count=0;
    while (getline(fin,line)) {
      if (line[0] == '#') continue;
      count++;
      if (count < datafrac) continue;
      count = 0;
      std::istringstream linein(line);
      double r,mapsq,mxsq,junk,gamsq;
      double sigmap,siggamsq;
      linein >>r>>mapsq>>sigmap>>mxsq>>junk>>junk>>junk>>junk>>junk>>gamsq>>siggamsq;
      r *= ARCMIN;

#ifdef BMODE
      mapsq += (BMODE)*mxsq;
#endif

#ifdef DILLUT
      mapsq *= DILLUTSQ;
      gamsq *= DILLUTSQ;
      //sigmap *= DILLUTSQ;
      //siggamsq *= DILLUTSQ;
#endif

      if (r > map_minr && r < map_maxr) {
	rmap[k].push_back(r);
	map[k].push_back(mapsq);
	//sigmap[k].push_back(sigmap);
      }
      if (r > var_minr && r < var_maxr) {
	rvar[k].push_back(r);
	var[k].push_back(gamsq);
	//sigvar[k].push_back(sigvar);
      }
    }
    dbg<<"sizes = "<<rmap[k].size()<<", "<<rvar[k].size()<<endl;
    if (k>0) {
      Assert(rmap[k].size() == rmap[0].size());
      Assert(rvar[k].size() == rvar[0].size());
    }
  }
  size_t nmap = rmap[0].size();
  size_t nvar = rvar[0].size();
  size_t ntot = nmap + nvar;
  dbg<<"nmap = "<<nmap<<", nvar = "<<nvar<<endl;

  r = new Vector<double>(ntot,0.);
  mean = new Vector<double>(ntot,0.);
  cov = new Matrix<double>(ntot,ntot,0.);

  for(size_t k=0;k<nfields;k++) {
    for(size_t i=0;i<nmap;i++) {
      (*r)[i] += rmap[k][i];
      (*mean)[i] += map[k][i];
      for(size_t j=0;j<=i;j++) {
	(*cov)(i,j) += map[k][i] * map[k][j];
      }
    }
    for(size_t i=0;i<nvar;i++) {
      (*r)[nmap+i] += rvar[k][i];
      (*mean)[nmap+i] += var[k][i];
      for(size_t j=0;j<=i;j++) {
	(*cov)(nmap+i,nmap+j) += var[k][i] * var[k][j];
      }
      for(size_t j=0;j<nmap;j++) {
	(*cov)(nmap+i,j) += var[k][i]*map[k][j];
      }
    }
  }
  for(size_t i=1;i<ntot;i++) cov->diag(i) = cov->diag(-i);

  *r /= double(nfields);
  *mean /= double(nfields);
  *cov -= double(nfields)*(*mean ^ *mean);
  *cov /= double(nfields*(nfields-1));
  for(size_t i=0;i<ntot;i++) Assert((*cov)(i,i) >= 0.);

#ifdef BMODE
  *cov /= sqrt(BMODE*BMODE+1);
#endif

  nmapout = nmap;

}

double PredInterp(double r, Vector<double> vlogr, Vector<double> vlogy)
{
  size_t n = vlogr.size();
  Assert(vlogy.size() == n);
  Assert(n >= 2);

  double logr1,logr2,logy1,logy2;
  double logr = log(r);
  if (logr < vlogr[0]) {
    dbg<<"Warning: extrapolating "<<r<<" from "<<exp(vlogr[0])<<", "<<exp(vlogr[1])<<endl;
    logy1 = vlogy[0];
    logy2 = vlogy[1];
    logr1 = vlogr[0];
    logr2 = vlogr[1];
  } else if (logr > vlogr[n-1]) {
    dbg<<"Warning: extrapolating "<<r<<" from "<<exp(vlogr[n-2])<<", "<<exp(vlogr[n-1])<<endl;
    logy1 = vlogy[n-2];
    logy2 = vlogy[n-1];
    logr1 = vlogr[n-2];
    logr2 = vlogr[n-1];
  } else {
    size_t i;
    for(i = 1; logr > vlogr[i]; i++) {
      Assert(i < n-1);
    }
    Assert(i < n);
    logy1 = vlogy[i-1];
    logy2 = vlogy[i];
    logr1 = vlogr[i-1];
    logr2 = vlogr[i];
  }

  double logy = ((logr-logr1)*logy2 + (logr2-logr)*logy1)/(logr2-logr1);
  return exp(logy);
}

int main(int argc, char* argv[])
{
  dbgout = &std::cout;
  double Om_de = -1.;

  // Best fit for w = -1
  double Om_m = 0.2575;
  double Om_b = 0.04275;
  double h = 0.7125;
  double As = 0.6675;
  double n = 0.9475;
  double w = -1.0;
  double wa = 0.0;

  /*
  // Best fit for w<-1 allowed
  double Om_m = 0.2625;
  double Om_b = 0.04525;
  double h = 0.7025;
  double As = 0.675;
  double n = 0.9425;
  double w = -0.9425;
  double wa = 0.0;
  */

  /*
  // Best fit for w,wa
  double Om_m = 0.429167;
  double Om_b = 0.0797917;
  double h = 0.6125;
  double As = 0.7875;
  double n = 0.964583;
  double w = -1.69583;
  double wa = 1.97917;
  */

  /*
  // Best fit for w,wa - not flat
  double Om_m = 0.315;
  Om_de = 0.675;
  double Om_b = 0.0505;
  double h = 0.685;
  double As = 0.785;
  double n = 0.9725;
  double w = -1.23;
  double wa = 1.23;
  */

  /*
  // Best fit for 0.95 dillution correction
  double Om_m = 0.23375;
  double Om_b = 0.04425;
  double h = 0.725;
  double As = 0.79;
  double n = 0.95875;
  double w = -0.9325;
  double wa = 0.0;
  */

  /*
  // Best fit for 1.05 dillution correction
  double Om_m = 0.26375;
  double Om_b = 0.04575;
  double h = 0.705;
  double As = 0.67;
  double n = 0.94375;
  double w = -0.9325;
  double wa = 0.0;
  */

  /*
  // Best fit for > 4 arcmin
  double Om_m = 0.23375;
  double Om_b = 0.04425;
  double h = 0.725;
  double As = 0.79;
  double n = 0.95875;
  double w = -0.9325;
  double wa = 0.0;
  */

  /*
  // Best fit for w<-1 not allowed
  double Om_m = 0.2675;
  double Om_b = 0.04925;
  double h = 0.675;
  double As = 0.78;
  double n = 0.9625;
  double w = -0.79;
  double wa = 0.0;
  */

  if (Om_de == -1.) Om_de = 1.-Om_m;

  if (argc == 2) As = atof(argv[1]);
  cout <<"As = "<<As<<endl;

  // Om_m  Om_DE  Om_b  h  A_s  n  w  wa
  const Universe u(Om_m,Om_de,Om_b,h,As,n,w,wa);

  Vector<double> pred_maplogr(npred_map+1);
  Vector<double> pred_logmap(npred_map+1);
  Vector<double> pred_varlogr(npred_var+1);
  Vector<double> pred_logvar(npred_var+1);
  double pred_map_logminr = log(pred_map_minr);
  double pred_var_logminr = log(pred_var_minr);
  double pred_map_step = exp(log(pred_map_maxr/pred_map_minr)/npred_map);
  double pred_var_step = exp(log(pred_var_maxr/pred_var_minr)/npred_var);
  double pred_map_logstep = log(pred_map_step);
  double pred_var_logstep = log(pred_var_step);

  std::ofstream mapout("mapsq.model");
  double r = pred_map_minr;
  for(int i=0;i<=npred_map;i++,r*=pred_map_step) {
    double mm = u.fmm(r);
    pred_maplogr(i) = log(r);
    pred_logmap(i) = log(mm);
    mapout << r/ARCMIN <<"  "<<mm<<endl;
  }

  std::ofstream varout("var.model");
  r = pred_var_minr;
  for(int i=0;i<=npred_var;i++,r*=pred_var_step) {
    double var = u.fvar(r);
    pred_varlogr(i) = log(r);
    pred_logvar(i) = log(var);
    varout << r/ARCMIN <<"  "<<var<<endl;
  }

  double sig8 = u.sigma8();
  cout <<"sigma8 = "<<sig8<<endl;

  Vector<double>* datar;
  Vector<double>* data;
  Matrix<double>* Sigma;
  size_t nmap;
  ReadData(datar,data,Sigma,nmap);
  size_t ntot = datar->size();

  Sigma->divideUsing(tmv::SV);
  Sigma->setDiv();
  Sigma->svd().thresh(singthresh,dbgout);

  Vector<double> pred(ntot);
  for(size_t i=0;i<ntot;i++) {
    if (i < nmap) {
      pred[i] = PredInterp((*datar)[i],pred_maplogr,pred_logmap);
      dbg<<"Map: "<<(*datar)[i]/ARCMIN<<"  "<<(*data)[i]<<"  "<<pred[i]<<endl;
    } else {
      pred[i] = PredInterp((*datar)[i],pred_varlogr,pred_logvar);
      dbg<<"Var: "<<(*datar)[i]/ARCMIN<<"  "<<(*data)[i]<<"  "<<pred[i]<<endl;
    }
  }

  Vector<double> diff = *data - pred;
  double chisq = diff * (diff/(*Sigma));
  cout <<"chisq = "<<chisq<<endl;
  cout <<"dof = "<<Sigma->svd().getKMax()-1<<endl;

  return 0;
}
